package com.training.player;

import android.net.Uri;

import androidx.annotation.OptIn;
import androidx.media3.common.util.UnstableApi;
import androidx.media3.datasource.BaseDataSource;
import androidx.media3.datasource.DataSource;
import androidx.media3.datasource.DataSpec;

import com.training.net.DataCenter;
import com.training.utils.AESUtil;
import java.io.ByteArrayOutputStream;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import javax.crypto.SecretKey;

/**
 * 自定义加密数据源，用于播放AES加密的音频文件
 * 
 * 此实现针对整体加密的文件，先将整个加密文件下载到内存，然后进行解密，
 * 最后按需提供解密后的数据块。
 */
@UnstableApi public class EncryptedDataSource extends BaseDataSource implements DataSource {
    private final String userAgent;
    private final SecretKey secretKey;
    private Uri uri;
    private byte[] decryptedData;
    private int currentPosition = 0;
    private boolean opened = false;
    private long bytesRemaining = 0;

    @OptIn(markerClass = UnstableApi.class) public EncryptedDataSource(String userAgent) {
        super(false);
        this.userAgent = userAgent;
        // 使用AESUtil中定义的密钥
//        this.secretKey = AESUtil.loadKey("s3cr3tK3yF0r4udi0");
        this.secretKey =AESUtil.loadKey(DataCenter.instance().aesKey) ;
    }

    @OptIn(markerClass = UnstableApi.class) @Override
    public long open(DataSpec dataSpec) throws IOException {
        try {
            this.uri = dataSpec.uri;
            
            // 如果已经打开过并且数据已解密，直接返回
            if (opened && decryptedData != null) {
                currentPosition = (int) dataSpec.position;
                bytesRemaining = decryptedData.length - currentPosition;
                transferStarted(dataSpec);
                return bytesRemaining;
            }
            
            // 下载并解密整个文件
            byte[] encryptedData = downloadEncryptedData(dataSpec.uri.toString());
            decryptedData = AESUtil.decrypt(encryptedData, secretKey);
            
            currentPosition = (int) dataSpec.position;
            bytesRemaining = decryptedData.length - currentPosition;
            
            opened = true;
            transferStarted(dataSpec);
            return bytesRemaining;
        } catch (Exception e) {
            throw new IOException("Failed to open encrypted data source", e);
        }
    }
    
    /**
     * 下载整个加密文件
     * @param urlString URL地址
     * @return 加密的数据
     * @throws IOException 网络异常
     */
    private byte[] downloadEncryptedData(String urlString) throws IOException {
        URL url = new URL(urlString);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        connection.setRequestProperty("User-Agent", userAgent);
        connection.connect();
        
        try (InputStream inputStream = connection.getInputStream();
             ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
            
            byte[] buffer = new byte[8192];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }
            
            return outputStream.toByteArray();
        } finally {
            connection.disconnect();
        }
    }

    @Override
    public int read(byte[] buffer, int offset, int readLength) throws IOException {
        if (readLength == 0) {
            return 0;
        } else if (bytesRemaining == 0) {
            return -1;
        }

        // 计算实际可读取的字节数
        int bytesToRead = (int) Math.min(readLength, bytesRemaining);
        
        // 从解密数据中复制到输出缓冲区
        System.arraycopy(decryptedData, currentPosition, buffer, offset, bytesToRead);
        
        // 更新位置和剩余字节数
        currentPosition += bytesToRead;
        bytesRemaining -= bytesToRead;
        
        return bytesToRead;
    }

    @Override
    public Uri getUri() {
        return uri;
    }

    @Override
    public void close() throws IOException {
        if (opened) {
            opened = false;
            decryptedData = null;
            currentPosition = 0;
            bytesRemaining = 0;
        }
    }
}